#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/applink.c>

#include <cppcodec/base64_rfc4648.hpp>

using base64 = cppcodec::base64_rfc4648;

// 错误处理函数
void handle_openssl_error() {
    char err_buf[512];
    ERR_error_string_n(ERR_get_error(), err_buf, sizeof(err_buf));
    fprintf(stderr, "OpenSSL error: %s\n", err_buf);
    exit(EXIT_FAILURE);
}

// 生成 RSA-2048 密钥对
void generate_rsa_keypair(const char* private_key_path, const char* public_key_path) {
    // 生成密钥
    RSA *rsa = RSA_generate_key(2048, RSA_F4, NULL, NULL);
    if (!rsa) {
        fprintf(stderr, "Failed to generate RSA key\n");
        handle_openssl_error();
    }

    // 打开文件
    FILE *pri_file = fopen(private_key_path, "w");
    FILE *pub_file = fopen(public_key_path, "w");
    if (!pri_file || !pub_file) {
        perror("Failed to open key files");
        exit(EXIT_FAILURE);
    }

    // 写入私钥
    if (!PEM_write_RSAPrivateKey(pri_file, rsa, NULL, NULL, 0, NULL, NULL)) {
        fprintf(stderr, "Failed to write private key\n");
        handle_openssl_error();
    }

    // 写入公钥
    if (!PEM_write_RSA_PUBKEY(pub_file, rsa)) {
        fprintf(stderr, "Failed to write public key\n");
        handle_openssl_error();
    }

    // 清理
    fclose(pri_file);
    fclose(pub_file);
    RSA_free(rsa);
}

// 对数据进行签名
int rsa_sign(const char* private_key_path, const unsigned char* data, size_t data_len,
             unsigned char** signature, size_t* sig_len) {
    FILE *fp = fopen(private_key_path, "r");
    if (!fp) {
        perror("Failed to open private key file");
        return 0;
    }

    EVP_PKEY *pkey = PEM_read_PrivateKey(fp, NULL, NULL, NULL);
    EVP_MD_CTX *ctx = EVP_MD_CTX_new();
        if (!ctx) return 0;

        if (EVP_DigestSignInit(ctx, NULL, EVP_sha256(), NULL, pkey) != 1) {
            EVP_MD_CTX_free(ctx);
            return 0;
        }

        if (EVP_DigestSignUpdate(ctx, data, data_len) != 1) {
            EVP_MD_CTX_free(ctx);
            return 0;
        }

        // 获取签名长度
        if (EVP_DigestSignFinal(ctx, NULL, sig_len) != 1) {
            EVP_MD_CTX_free(ctx);
            return 0;
        }

        *signature = (unsigned char *)OPENSSL_malloc(*sig_len);
        if (!*signature) {
            EVP_MD_CTX_free(ctx);
            return 0;
        }

        if (EVP_DigestSignFinal(ctx, *signature, sig_len) != 1) {
            OPENSSL_free(*signature);
            EVP_MD_CTX_free(ctx);
            return 0;
        }

        EVP_MD_CTX_free(ctx);
        EVP_PKEY_free(pkey);

    return 1;
}

// 验证签名
int rsa_verify(const char* public_key_path, const unsigned char* data, size_t data_len,
               const unsigned char* signature, size_t sig_len) {
    FILE *fp = fopen(public_key_path, "r");
    if (!fp) {
        perror("Failed to open public key file");
        return 0;
    }

#if 0
    // 读取公钥
    RSA *rsa = PEM_read_RSA_PUBKEY(fp, NULL, NULL, NULL);
    fclose(fp);

    if (!rsa) {
        fprintf(stderr, "Failed to read public key\n");
        return 0;
    }

    // 创建消息摘要上下文
    EVP_MD_CTX *ctx = EVP_MD_CTX_new();
    if (!ctx) {
        perror("Failed to create EVP context");
        RSA_free(rsa);
        return 0;
    }

    // 初始化验证操作 - 使用 SHA-256 算法
    int result = EVP_DigestVerifyInit(ctx, NULL, EVP_sha256(), NULL, NULL) &&
                 EVP_DigestVerifyUpdate(ctx, data, data_len) &&
                 (EVP_DigestVerifyFinal(ctx, signature, sig_len) == 1);

    // 清理
    EVP_MD_CTX_free(ctx);
    RSA_free(rsa);
#endif
    // 读取公钥
    EVP_PKEY *pkey = PEM_read_PUBKEY(fp, NULL, NULL, NULL);//   PEM_read_PrivateKey(fp, NULL, NULL, NULL);
    fclose(fp);

    if (!pkey) {
        fprintf(stderr, "Failed to read public key\n");
        return 0;
    }

    // 创建消息摘要上下文
    EVP_MD_CTX *ctx = EVP_MD_CTX_new();
    if (!ctx) {
        perror("Failed to create EVP context");
        EVP_PKEY_free(pkey);
        return 0;
    }

    // 初始化验证操作 - 使用 SHA-256 算法
    int result = EVP_DigestVerifyInit(ctx, NULL, EVP_sha256(), NULL, pkey) &&
                 EVP_DigestVerifyUpdate(ctx, data, data_len) &&
                 (EVP_DigestVerifyFinal(ctx, signature, sig_len) == 1);

    // 清理
    EVP_MD_CTX_free(ctx);
    EVP_PKEY_free(pkey);
    return result;
}

#define LOGW(format, ...) printf(format "\n", ##__VA_ARGS__)

EVP_PKEY* create_evp_pkey_from_pem_string(const char* pem_str, int is_public) {
    BIO* bio = BIO_new_mem_buf(pem_str, -1);
    if (!bio) {
        LOGW("Error creating BIO buffer");
        return NULL;
    }

    EVP_PKEY* pkey = NULL;
    if (is_public) {
        pkey = PEM_read_bio_PUBKEY(bio, NULL, NULL, NULL);
    } else {
        pkey = PEM_read_bio_PrivateKey(bio, NULL, NULL, NULL);
    }

    BIO_free(bio);

    if (!pkey) {
        LOGW("Error loading key: %s", ERR_error_string(ERR_get_error(), NULL));
    }

    return pkey;
}

// 验证签名
int rsa_verify2(const char* pem_str, const unsigned char* data, size_t data_len,
               const unsigned char* signature, size_t sig_len) {
//    FILE *fp = fopen(public_key_path, "r");
//    if (!fp) {
//        perror("Failed to open public key file");
//        return 0;
//    }
//    // 读取公钥
//    EVP_PKEY *pkey = PEM_read_PUBKEY(fp, NULL, NULL, NULL);//   PEM_read_PrivateKey(fp, NULL, NULL, NULL);
//    fclose(fp);
    EVP_PKEY *pkey = create_evp_pkey_from_pem_string(pem_str, 1);
    if (!pkey) {
        LOGW("Failed to read public key");
        return 0;
    }

    // 创建消息摘要上下文
    EVP_MD_CTX *ctx = EVP_MD_CTX_new();
    if (!ctx) {
        LOGW("Failed to create EVP context");
        EVP_PKEY_free(pkey);
        return 0;
    }

    // 初始化验证操作 - 使用 SHA-256 算法
    int result = EVP_DigestVerifyInit(ctx, NULL, EVP_sha256(), NULL, pkey) &&
                 EVP_DigestVerifyUpdate(ctx, data, data_len) &&
                 (EVP_DigestVerifyFinal(ctx, signature, sig_len) == 1);
    if (!result) {
        handle_openssl_error();
    }

    // 清理
    EVP_MD_CTX_free(ctx);
    EVP_PKEY_free(pkey);
    return result;
}
int main() {
    // 初始化 OpenSSL
    OpenSSL_add_all_algorithms();
    ERR_load_crypto_strings();

    const char* private_key_path = "./keys/private.pem";
    const char* public_key_path = "./keys/public.pem";

    // 生成RSA-2048密钥对
    // printf("Generating RSA-2048 key pair...\n");
    // generate_rsa_keypair(private_key_path, public_key_path);
    // printf("Key pair generated\n");

    // 测试数据
    const char* message = "/osc/info/1749519966/98989898";
    size_t msg_len = strlen(message);

    // 签名
    unsigned char* signature = NULL;
    size_t sig_len;

    printf("\nSigning message: '%s'\n", message);
    if (!rsa_sign(private_key_path, (const unsigned char*)message, msg_len, &signature, &sig_len)) {
        fprintf(stderr, "Signature failed\n");
        return 1;
    }
    char buf[2048] = {0};
    size_t size = base64::encode(buf, 2048, (const char *)signature, sig_len);
    printf("Signature generated (length: %zu)\n%s\n", sig_len, (char *)buf);


    static const char* pem_str =
        "-----BEGIN PUBLIC KEY-----\n"
        "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA02GadEus6dv4HCuZHdm+\n"
        "EytbvW4v+BJV1FOXFGr3WyB2+dpCkwTHosQPv48bKy4t7Vw5cqfv89UehA7m7Ubt\n"
        "htvd1cXmYbcYlSGy9B7J5DxOWuWiTN/vID7dj1OuplLfwjlURH7kiXX7dEPAhEhA\n"
        "055lGNUOVljX+1YqDE2+URRhjFu2BkYadQjrUnRlP3YQQuWW5aIa0sqkKUK5WAXl\n"
        "vlXLT6omwl/YT7xCrh03F7tPRL/J6+Q380qRYydPKyT244gPlYm4Jngj7ZxXpAod\n"
        "jeI+mAPusKwXSXt6OhvT+2mmyEEJc589Cl2Ai+ugxjdAOxeXdYLBMpbspzEqvsva\n"
        "kwIDAQAB\n"
        "-----END PUBLIC KEY-----\n";

    // 验证签名（预期成功）
    printf("\nVerifying signature...\n");
    if (rsa_verify2(pem_str, (const unsigned char*)message, msg_len, signature, sig_len)) {
        printf("Signature verification SUCCESS\n");
    } else {
        printf("Signature verification FAILED (expected success)\n");
    }

    // 篡改签名并再次验证（预期失败）
    printf("\nTesting with modified signature...\n");
    signature[0] ^= 0xFF;  // 篡改签名内容
    if (rsa_verify(public_key_path, (const unsigned char*)message, msg_len, signature, sig_len)) {
        printf("Signature verification SUCCESS (but should have failed!)\n");
    } else {
        printf("Signature verification FAILED (as expected)\n");
    }

    // 篡改消息并再次验证（预期失败）
    printf("\nTesting with modified message...\n");
    char modified_msg[1024];
    snprintf(modified_msg, sizeof(modified_msg), "%s", message);
    modified_msg[0] ^= 0xFF;  // 篡改消息内容

    if (rsa_verify(public_key_path, (const unsigned char*)modified_msg, strlen(modified_msg),
                   signature, sig_len)) {
        printf("Signature verification SUCCESS (but should have failed!)\n");
    } else {
        printf("Signature verification FAILED (as expected)\n");
    }

    // 清理
    OPENSSL_free(signature);

    // 清理OpenSSL资源
    EVP_cleanup();
    CRYPTO_cleanup_all_ex_data();
    ERR_free_strings();

    return 0;
}
